package com.example.mcp.server.service.handler;

/**
 * @author ssk
 * @date 2025/8/8
 **/

import com.example.mcp.server.model.McpInitializeRequest;
import com.example.mcp.server.model.McpInitializeResponse;
import com.example.mcp.server.model.McpListToolsRequest;
import com.example.mcp.server.model.McpListToolsResponse;
import com.example.mcp.server.model.McpToolCallRequest;
import com.example.mcp.server.model.McpToolCallResponse;
import com.example.mcp.server.service.McpService;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

@Component
public class McpWebSocketHandler extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(McpWebSocketHandler.class);

    private final ObjectMapper objectMapper = new ObjectMapper();
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    @Autowired
    private McpService mcpService;


    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        logger.debug("收到 WebSocket 消息: {}", payload);

        try {
            // 解析消息类型
            Map<String, Object> messageMap = objectMapper.readValue(payload, Map.class);
            String type = (String) messageMap.get("type");

            switch (type) {
                case "initialize":
                    handleInitialize(session, payload);
                    break;
                case "tools/call":
                    handleToolCall(session, payload);
                    break;
                case "tools/list":
                    handleListTools(session, payload);
                    break;
                default:
                    sendError(session, "Unknown message type: " + type);
            }
        } catch (Exception e) {
            logger.error("处理WebSocket消息时发生错误: {}", e.getMessage(), e);
            sendError(session, "Error processing message: " + e.getMessage());
        }
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = UUID.randomUUID().toString();
        sessions.put(sessionId, session);
        session.getAttributes().put("sessionId", sessionId);
        logger.info("WebSocket 连接已建立，会话ID: {}", sessionId);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        String sessionId = (String) session.getAttributes().get("sessionId");
        if (sessionId != null) {
            sessions.remove(sessionId);
            mcpService.closeSession(sessionId);
        }
        logger.info("WebSocket 连接已关闭，会话ID: {}，关闭状态: {}", sessionId, status);
    }

    private void handleInitialize(WebSocketSession session, String payload) throws Exception {
        McpInitializeRequest request = objectMapper.readValue(payload, McpInitializeRequest.class);
        String sessionId = (String) session.getAttributes().get("sessionId");

        McpInitializeResponse response = mcpService.initialize(request);
        response.setId(request.getId());

        sendMessage(session, objectMapper.writeValueAsString(response));
    }

    private void handleToolCall(WebSocketSession session, String payload) throws Exception {
        McpToolCallRequest request = objectMapper.readValue(payload, McpToolCallRequest.class);
        String sessionId = (String) session.getAttributes().get("sessionId");

        if (!mcpService.validateSession(sessionId)) {
            sendError(session, "Invalid session");
            return;
        }

        McpToolCallResponse response = mcpService.callTool(request);
        response.setId(request.getId());

        sendMessage(session, objectMapper.writeValueAsString(response));
    }

    private void handleListTools(WebSocketSession session, String payload) throws Exception {
        McpListToolsRequest request = objectMapper.readValue(payload, McpListToolsRequest.class);
        String sessionId = (String) session.getAttributes().get("sessionId");

        if (!mcpService.validateSession(sessionId)) {
            sendError(session, "Invalid session");
            return;
        }

        McpListToolsResponse response = mcpService.listTools(sessionId);
        response.setId(request.getId());

        sendMessage(session, objectMapper.writeValueAsString(response));
    }

    private void sendError(WebSocketSession session, String error) {
        Map<String, Object> errorResponse = Map.of(
                "type", "error",
                "error", error
        );

        try {
            sendMessage(session, objectMapper.writeValueAsString(errorResponse));
        } catch (Exception e) {
            logger.error("发送错误消息失败: {}", e.getMessage(), e);
        }
    }

    private void sendMessage(WebSocketSession session, String message) {
        try {
            if (session.isOpen()) {
                session.sendMessage(new TextMessage(message));
            } else {
                logger.warn("WebSocket会话已关闭，无法发送消息");
                // 清理已关闭的会话
                cleanupClosedSession(session);
            }
        } catch (Exception e) {
            logger.error("发送WebSocket消息时发生错误: {}", e.getMessage(), e);
            // 检查是否是连接重置异常
            if (isConnectionReset(e)) {
                logger.info("检测到连接重置，清理会话");
                cleanupClosedSession(session);
            } else {
                // 重新抛出其他异常
                throw new RuntimeException("发送消息失败", e);
            }
        }
    }

    private boolean isConnectionReset(Exception e) {
        return e instanceof java.io.IOException &&
                (e.getMessage().contains("Connection reset by peer") ||
                        e.getMessage().contains("Broken pipe"));
    }

    private void cleanupClosedSession(WebSocketSession session) {
        String sessionId = (String) session.getAttributes().get("sessionId");
        if (sessionId != null) {
            sessions.remove(sessionId);
            try {
                mcpService.closeSession(sessionId);
            } catch (Exception e) {
                logger.error("关闭会话时发生错误: {}", e.getMessage(), e);
            }
        }
    }
}